import os
import sys
import json
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt
from sklearn.metrics import classification_report, roc_curve, auc
from sklearn.model_selection import train_test_split
from tensorflow.keras.utils import to_categorical


# 把项目根目录加入 sys.path
ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
sys.path.insert(0, ROOT)

# 现在可以按包路径引用 models 里的模块
from models.model_logistic    import LogisticRegressionModel
from models.model_mlp         import MLP
from models.model_rf          import get_rf_model
from models.model_lgbm        import get_lgbm_model
from models.model_xgb         import get_xgb_model
from models.model_cnn         import CNNModel
from models.model_lstm        import LSTMModel
from models.model_transformer import TransformerModel
from models.model_catboost    import CatBoostModel
# Safe ROC calculation
def safe_roc(y_true, y_score):
    mask = (y_score is not None) and ~np.isnan(y_score)
    if y_score is None:
        return None, None, None
    y_true_f = y_true[mask]
    y_score_f = y_score[mask]
    if len(np.unique(y_true_f)) < 2:
        return None, None, None
    try:
        fpr, tpr, _ = roc_curve(y_true_f, y_score_f)
        return fpr, tpr, auc(fpr, tpr)
    except:
        return None, None, None

# Model factory
def get_model(name, input_dim=None, output_dim=None):
    if name == 'Logistic':
        return LogisticRegressionModel(input_dim=input_dim, output_dim=output_dim)
    if name == 'MLP':
        return MLP(input_dim=input_dim, hidden_dim=64, output_dim=output_dim)
    if name == 'RandomForest':
        return get_rf_model()
    if name == 'XGBoost':
        return get_xgb_model()
    if name == 'LightGBM':
        return get_lgbm_model()
    if name == 'CNN':
        return CNNModel(input_shape=(None,1), num_classes=output_dim)
    if name == 'LSTM':
        return LSTMModel(input_shape=(None,1), num_classes=output_dim)
    if name == 'Transformer':
        return TransformerModel(input_dim=1, num_classes=output_dim)
    if name == 'CatBoost':
        return CatBoostModel()
    raise ValueError(f"Unknown model: {name}")

# Train and evaluate a model
def evaluate_model(name, model, X_tr, y_tr, X_te, y_te,
                   report_dir, pic_dir, roc_dir, task):
    print(f"\n=== Task: {task}, Model: {name} ===")
    # Train
    if hasattr(model, 'train'):
        model.train(X_tr, y_tr)
    elif hasattr(model, 'fit'):
        model.fit(X_tr, y_tr)
    # Predict
    if hasattr(model, 'predict_proba'):
        probs = model.predict_proba(X_te)
        score = probs[:,1] if probs.ndim>1 else probs
        y_pred = model.predict(X_te)
    elif hasattr(model, 'predict') and not isinstance(model, torch.nn.Module):
        y_pred = model.predict(X_te)
        score = None
    else:
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        model.to(device)
        Xte_t = torch.FloatTensor(X_te).to(device)
        with torch.no_grad():
            logits = model(Xte_t)
        if logits.ndim == 2 and logits.shape[1] > 1:
            probs = torch.softmax(logits, dim=1).cpu().numpy()
            score = None
            y_pred = probs.argmax(axis=1)
        else:
            probs = torch.sigmoid(logits.view(-1)).cpu().numpy()
            score = probs
            y_pred = (probs>0.5).astype(int)
    # Print raw predictions
    print("Predictions:", y_pred)
    if score is not None:
        print("Scores:", score)
    # Classification report
    rpt = classification_report(y_te, y_pred, zero_division=0)
    print(rpt)
    pd.DataFrame(classification_report(y_te,y_pred,output_dict=True,zero_division=0)).T.to_csv(
        os.path.join(report_dir, f"{task}_{name}.csv")
    )
    # Feature importance
    if hasattr(model, 'feature_importances_'):
        plot_feature_importance(model.feature_importances_, None,
                                os.path.join(pic_dir, f"{task}_{name}.png"),
                                title=f"{task}_{name} Feature Importance")
    # ROC curve
    fpr, tpr, roc_auc = safe_roc(y_te, score)
    plt.figure()
    if fpr is not None:
        plt.plot(fpr, tpr, label=f'AUC={roc_auc:.2f}')
        plt.plot([0,1],[0,1],'k--')
        plt.legend(loc='lower right')
    plt.title(f'{task}_{name}_ROC')
    plt.xlabel('FPR'); plt.ylabel('TPR')
    plt.savefig(os.path.join(roc_dir, f"{task}_{name}_roc.png"))
    plt.close()

if __name__ == '__main__':
    # Load config
    cfg = json.load(open('config.json'))
    sel_models = cfg.get('models', [])
    sel_tasks  = cfg.get('tasks', [])
    # Load data
    df = pd.read_csv(r"C:\py\hrv_predict\data_hrv\hrv_dataset.csv")
    df['age_group'] = (df['age']//10).astype(int)
    df['gender']    = df['Gender'].map({'男':1,'女':0}).fillna(0).astype(int)
    X = df.drop(columns=['record','age','Gender','AnalysisText','age_group','gender']).values
    y_age    = df['age_group'].values
    y_gender = df['gender'].values
    # Split
    Xa_tr,Xa_te,ya_tr,ya_te = train_test_split(X,y_age,test_size=0.2,random_state=0,stratify=y_age)
    Xg_tr,Xg_te,yg_tr,yg_te = train_test_split(X,y_gender,test_size=0.2,random_state=0,stratify=y_gender)
    # Output dirs
    rpt_dir = r"C:\py\hrv_predict\result\report"; pic_dir = r"C:\py\hrv_predict\result\pic"; roc_dir = r"C:\py\hrv_predict\result\roc"
    for d in (rpt_dir,pic_dir,roc_dir): os.makedirs(d, exist_ok=True)
    # Run
    for task in ['age','gender']:
        if task not in sel_tasks: continue
        Xtr,Xte,ytr,yte = (Xa_tr,Xa_te,ya_tr,ya_te) if task=='age' else (Xg_tr,Xg_te,yg_tr,yg_te)
        out_dim = len(np.unique(ytr)) if task=='age' else 1
        for m in sel_models:
            try:
                model = get_model(m, input_dim=X.shape[1], output_dim=out_dim)
                evaluate_model(m, model, Xtr, ytr, Xte, yte, rpt_dir, pic_dir, roc_dir, task)
            except Exception as e:
                print(f"Skipping model {m} due to error: {e}")
